from construct2.protocols.ipstack import ip_stack


def sniffer():
    """use Ctrl+C to stop"""
    import pcap
    for ts, pkt in pcap.pcap(promisc = False):
        yield ip_stack.parse(pkt)

def tcp_follower(packets, addr1, port1, addr2, port2):
    """follows a packets that belong to the given tcp stream:
    (source address, source port, destination address, destination port)
    
    this function yields each matching packet as the tuple:
    ((seq, ack), (side, data)).
    
    when coming across a FIN packet, the stream is considered "done with"
    """
    for p in packets:
        if p.header.type != "IPv4":
            continue
        if p.next.header.protocol != "TCP":
            continue
        if not p.next.next.next.data:
            continue
        sa = p.next.header.source
        sp = p.next.next.header.source
        da = p.next.header.destination
        dp = p.next.next.header.destination
        key = (p.next.next.header.seq, p.next.next.header.ack)
        if sa == addr1 and sp == port1 and da == addr2 and dp == port2:
            yield key, (1, p.next.next.next.data)
            if p.next.next.header.flags.fin:
                break
        elif sa == addr2 and sp == port2 and da == addr1 and dp == port1:
            yield key, (2, p.next.next.next.data)
            if p.next.next.header.flags.fin:
                break

def tcp_stream(follower):
    """a trivial sequencer of tcp packets. returns a tuple of the extracted,
    sequenced, data of each side as (side_a, side_b). 
    it should actually implement a sliding-window instead of holding the
    entire session in-memory, but that's for a later implementation
    """
    side_a = {}
    side_b = {}
    for key, (side, value) in follower:
        if side == 1:
            side_a[key] = value
        else:
            side_b[key] = value
    side_a_data = "".join(v for k, v in sorted(side_a.items()))
    side_b_data = "".join(v for k, v in sorted(side_b.items()))
    return side_a_data, side_b_data

def tcp_streamify(packets):
    streams = {}
    for p in packets:
        if p.header.type != "IPv4":
            continue
        if p.next.header.protocol != "TCP":
            continue
        if not p.next.next.next.data:
            continue
        sa = p.next.header.source
        sp = p.next.next.header.source
        da = p.next.header.destination
        dp = p.next.next.header.destination
        sidea = (sa, sp)
        sideb = (da, dp)
        side = sidea is max(sidea, sideb)
        sid = frozenset((sidea, sideb))
        stream = streams.setdefault(sid, {})
        pktid = p.next.next.header.seq, p.next.next.header.ack
        stream[side, pktid] = p.next.next.next.data
    for k, stream in streams.iteritems():
        items = stream.items()
        items.sort()
        sides = {0 : [], 1 : []}
        for (side, pktid), value in items:
            sides[side].append(value)
        streams[k] = ("".join(sides[0]), "".join(sides[1]))
    return streams


if __name__ == "__main__":
    import cPickle, itertools
    pkts = cPickle.load(open("cap.pkl", "rb"))
    ipkts = (ip_stack.parse(p) for p in pkts)
    
    for k, (sidea, sideb) in tcp_streamify(ipkts).iteritems():
        print sidea
        print "-" * 80
        print sideb
        print
        break
    
#    follower = tcp_follower(ipkts, '192.168.2.4', 3156, '82.94.237.218', 80)
#    side_a, side_b = tcp_stream(follower)
#    print side_a
#    print "-" * 80
#    print side_b

    











